from strategy.core.Stg import Stg
import talib as ta
import math
import pandas as pd


class KDJStg(Stg):
    preBuySell = 0
    param = {
        'period': 9,
        'smooth': 3,
    }

    def initial(self):
        df = self.data
        kdj = pd.DataFrame()
        # KDJ 值对应的函数是 STOCH
        kdj['k'], kdj['d'] = ta.STOCH(
            df['high'].values,
            df['low'].values,
            df['close'].values,
            fastk_period=self.param['period'],
            slowk_period=self.param['smooth'],
            slowk_matype=0,
            slowd_period=self.param['smooth'],
            slowd_matype=0)
        # 求出J值，J = (3*K)-(2*D)
        kdj['j'] = list(map(lambda x, y: 3 * x - 2 * y, kdj['k'], kdj['d']))
        self.kdj = kdj
        if self.param['period'] == 5.1:
            self.highPoint = 90
        # 格式: self.macd结果包括: macd,macdsignal,macdhist. 分别对应dif,dea,dif-dea
        # self.fastSma = ta.SMA(self.data, self.param['fastPeriod'])\
        return

    def signal(self, index):
        # 算法
        # k或d<20: 超卖, 应买入
        # k或d>80: 超买, 应卖出
        # 自动跳过无法计算的节点
        if index < self.param['period'] + self.param['smooth']:
            return
        pre = index - 1
        res = 0

        if self.param['buySell'] == 'over':  # 超买超卖
            # 超卖, 买入
            if self.kdj['k'][index] < 100 - self.param['superBuyPoint'] or self.kdj['d'][index] < 100 - self.param[
                'superBuyPoint']:
                return self.buy()
            # 超买, 卖出, superBuyPoint一般是80
            if self.kdj['k'][index] > self.param['superBuyPoint'] or self.kdj['d'][index] > self.param['superBuyPoint']:
                return self.sell()
        if self.param['buySell'] == 'x':  # 金叉死叉
            # 金叉: k向上超越d, 买入
            if self.kdj['k'][pre] < self.kdj['k'][index] and self.kdj['k'][index] > self.kdj['d'][index]:
                return self.buy()
            # 死叉: k向下穿越d, 卖出
            if self.kdj['k'][pre] > self.kdj['k'][index] and self.kdj['k'][index] < self.kdj['d'][index]:
                return self.sell()

    def buy(self):  # 在上次买后不再买, 防止重复出现买状态
        if self.preBuySell != 1:
            self.preBuySell = 1
            return 1

    def sell(self):
        if self.preBuySell != -1:
            self.preBuySell = -1
            return -1
